package com.example.common.webSocket;

import org.json.JSONObject;
import org.springframework.stereotype.Component;

import javax.websocket.OnClose;
import javax.websocket.OnMessage;
import javax.websocket.OnOpen;
import javax.websocket.Session;
import javax.websocket.server.ServerEndpoint;
import java.util.concurrent.ConcurrentHashMap;

@Component
@ServerEndpoint("/progress")
public class ProgressWebSocket {
    private static final ConcurrentHashMap<String, Session> flagSessions = new ConcurrentHashMap<>();
    private static final ConcurrentHashMap<String, String> sessionFlags = new ConcurrentHashMap<>();

    @OnOpen
    public void onOpen(Session session) {
        System.out.println("WebSocket connected: " + session.getId());
    }

    @OnMessage
    public void onMessage(String message, Session session) {
        try {
            JSONObject json = new JSONObject(message);
            if ("subscribe".equals(json.getString("action"))) {
                String flag = json.getString("flag");
                // 打印日志验证
                System.out.println("收到订阅请求 - flag: " + flag);
                flagSessions.put(flag, session);
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    @OnClose
    public void onClose(Session session) {
        String flag = sessionFlags.get(session.getId());
        if (flag != null) {
            flagSessions.remove(flag);
            sessionFlags.remove(session.getId());
        }
        System.out.println("WebSocket closed: " + session.getId());
    }

    public static void sendProgress(String flag, int progress) {
        Session session = flagSessions.get(flag);
        if (session != null && session.isOpen()) {
            try {
                JSONObject json = new JSONObject();
                json.put("flag", flag);
                json.put("progress", progress);
                session.getBasicRemote().sendText(json.toString());
            } catch (Exception e) {
                e.printStackTrace();
            }
        }
    }
}